import torch
import numpy as np
import torch.nn as nn
import tool



class TripleLoss2(nn.Module):
    def __init__(self, weight=10):
        super(TripleLoss2, self).__init__()
        self.weight = weight

    def forward(self, ori, option_ori, cor_angle, cor_index):
        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = int(cor_index)

        distance_mac = 2 - 2 * torch.einsum("j,ijk->ik", [ori, option_ori])
        size = distance_mac.shape
        pos = distance_mac[cor_x][cor_y]
        triple_dis = pos - distance_mac
        loss = torch.sum(torch.log(1 + torch.exp(triple_dis * self.weight)))
        return loss / (size[0] * size[1] - 1)


class TripleLoss3(nn.Module):
    def __init__(self, weight=10):
        super(TripleLoss3, self).__init__()
        self.weight = weight
        self.semi_weight = 0.8 * weight

    def forward(self, ori, option_ori, cor_angle, cor_index):
        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = int(cor_index)

        distance_mac = 2 - 2 * torch.einsum("j,ijk->ik", [ori, option_ori])
        size = distance_mac.shape
        pos = distance_mac[cor_x][cor_y]
        triple_semi_dis = pos - distance_mac[cor_x, :]
        triple_hard_dis1 = pos - distance_mac[:cor_x, :]
        triple_hard_dis2 = pos - distance_mac[cor_x + 1:, :]
        triple_hard_dis = torch.cat([triple_hard_dis1, triple_hard_dis2], dim=0)
        loss = torch.sum(torch.log(1 + torch.exp(triple_hard_dis * self.weight))) + torch.sum(
            torch.log(1 + torch.exp(triple_semi_dis * self.semi_weight)))
        return loss / (size[0] * size[1] - 1)


class TripleLoss(nn.Module):
    def __init__(self, margin=0.6):
        super(TripleLoss, self).__init__()
        self.margin = margin
        self.semi_margin = 0.6

    def compute_loss_semi(self, pos, neg):
        loss = pos - neg + self.semi_margin
        return loss if loss > 0 else 0

    def compute_loss(self, pos, neg):
        loss = pos - neg + self.margin
        return loss if loss > 0 else 0

    def forward(self, ori, option_ori, cor_angle, cor_index):
        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = int(cor_index)

        distance_mac = torch.einsum("j,ijk->ik", [ori, option_ori])
        # distance_mac = 2 - 2 * torch.einsum("j,ijk->ik", [ori, option_ori])
        size = distance_mac.shape
        pos_dist = distance_mac[cor_x][cor_y]

        loss = 0
        # for i in range(size[0]):
        #     for j in range(size[1]):
        #         loss += self.compute_loss(pos_dist, distance_mac[i][j])
        #

        for i in range(size[0]):
            for j in range(size[1]):
                if i == cor_x:
                    if j != cor_y:
                        loss += self.compute_loss_semi(pos_dist, distance_mac[i][j])
                else:
                    loss += self.compute_loss(pos_dist, distance_mac[i][j])

        return loss / (size[0] * size[1] - 1)


class TripleLoss1(nn.Module):
    def __init__(self, margin=4):
        super(TripleLoss1, self).__init__()
        self.margin = margin
        self.semi_margin = 3.5

    def compute_loss_semi(self, pos, neg):
        loss = pos - neg + self.semi_margin
        return loss if loss > 0 else 0

    def compute_loss(self, pos, neg):
        loss = pos - neg + self.margin
        return loss if loss > 0 else 0

    def forward(self, ori, option_ori, cor_angle, cor_index):
        angle_list = np.arange(0, 360, 360.0 / option_ori.shape[-1])
        cor_y = tool.getrightlori(angle_list, cor_angle)
        cor_x = int(cor_index)

        distance_mac = 2 - 2 * torch.einsum("j,ijk->ik", [ori, option_ori])
        size = distance_mac.shape
        pos_dist = distance_mac[cor_x][cor_y]

        loss = 0
        # for i in range(size[0]):
        #     for j in range(size[1]):
        #         loss += self.compute_loss(pos_dist, distance_mac[i][j])
        #

        for i in range(size[0]):
            for j in range(size[1]):
                if i == cor_x:
                    if j != cor_y:
                        loss += self.compute_loss_semi(pos_dist, distance_mac[i][j])
                else:
                    loss += self.compute_loss(pos_dist, distance_mac[i][j])

        return loss / (size[0] * size[1] - 1)


# test
if __name__ == '__main__':
    pass
    # loss = TripleLoss3(10)
    # distance_mac =  torch.rand((4, 4))
    # for i in range(4):
    #     for j in range(4):
    #         distance_mac[i][j] = 4 * i + j
    # size = distance_mac.shape
    # cor_x = 3
    # cor_y = 1
    #
    # pos = distance_mac[cor_x][cor_y]
    # triple_semi_dis = pos - distance_mac[cor_x, :]
    # triple_hard_dis1 = pos - distance_mac[:cor_x, :]
    # triple_hard_dis2 = pos - distance_mac[cor_x+1:, :]
    # triple_hard_dis = torch.cat([triple_hard_dis1, triple_hard_dis2], dim=0)
    # print(distance_mac)
    # print(pos)
    # print(triple_semi_dis)
    # print(triple_hard_dis1)
    # print(triple_hard_dis2)
    # print(triple_hard_dis)
    # print('0')

